import os, argparse, json, platform, shutil, socket, subprocess, sys, tarfile, tempfile, time
from pathlib import Path
from urllib.parse import parse_qs
from urllib.request import Request, urlopen


DEFAULT_SERVER = "https://server-agent.gitee.com/gitee_sa_server"
DEFAULT_FILE_DOWNLOAD_URI = "https://gitee-agent-shell.gz.bcebos.com"
AGENT_DOWNLOAD_PATH = "/agent/agent.jar"
UUID_AUTH_URL_PATH = "/sa/rest/v2/agents"
JDK_TARBALL_BY_ARCH = {
    "x86_64": "/agent/jdk-8u251-linux-x64.tar.gz",
    "aarch64": "/agent/arm64/jdk-8u202-linux-arm64-vfp-hflt.tar.gz",
}
PROXY_PATH_BY_ARCH = {
    "x86_64": "/proxy/gitee-proxy",
    "aarch64": "/proxy/arm64/gitee-proxy",
}


def which(cmd: str) -> str | None:
    return shutil.which(cmd)


def is_arch_supported() -> tuple[str, str, str]:
    arch = platform.machine()
    if arch not in JDK_TARBALL_BY_ARCH:
        print(f"不支持的架构: {arch}，仅支持x86_64、arm64", file=sys.stderr)
        sys.exit(1)
    return arch, JDK_TARBALL_BY_ARCH[arch], PROXY_PATH_BY_ARCH[arch]


def ensure_dir(p: Path) -> None:
    p.mkdir(parents=True, exist_ok=True)


def http_get(url: str, *, headers: dict[str, str] | None = None, timeout: int = 30) -> bytes:
    req = Request(url, headers=headers or {})
    with urlopen(req, timeout=timeout) as resp:
        return resp.read()


def http_download(url: str, dst: Path, *, headers: dict[str, str] | None = None, timeout: int = 300) -> None:
    req = Request(url, headers=headers or {})
    with urlopen(req, timeout=timeout) as resp, open(dst, "wb") as f:
        shutil.copyfileobj(resp, f)


def monitor_ok(server: str) -> bool:
    url = f"{server.rstrip('/')}/monitor/ok"
    try:
        # 与原脚本一致：200/301/异常（模拟000）都认为可接受
        req = Request(url)
        with urlopen(req, timeout=30) as resp:
            return resp.status in (200, 301)
    except Exception:
        # 视作 000
        return True


def extract_tar_gz(src: Path, dst_dir: Path) -> None:
    with tarfile.open(src, mode="r:gz") as tar:
        tar.extractall(path=dst_dir)


def detect_java_version(java_cmd: str) -> str | None:
    try:
        # java -version 输出在 stderr
        proc = subprocess.run([java_cmd, "-version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        out = proc.stderr.splitlines() or proc.stdout.splitlines()
        if not out:
            return None
        line = out[0]
        # 常见格式：openjdk version "1.8.0_252" 或 "17.0.10"
        # 取第一个形如 "X.Y.Z" 的版本串
        import re
        m = re.search(r"(\d+)(?:\.(\d+))?(?:\.(\d+))?", line)
        return m.group(0) if m else None
    except Exception:
        return None


def version_second_bit_ge_6(ver: str) -> bool:
    # 取第二段（按 . 分割）并判断 >= 6（即 1.6+）
    # 1.8.0_251 -> second_bit=8；17.0.10 -> second_bit=0（视为 >=6? 否）
    try:
        parts = ver.split(".")
        if len(parts) >= 2:
            second = int(parts[1])
            return second >= 6
        # 没有第二段时，保守视为不满足
        return False
    except Exception:
        return False


def install_jdk(install_dir: Path, file_download_uri: str, server: str, label_type: str, jdk_path: str) -> Path:
    jdk_dir = install_dir / "jdk4agent"
    if not os.path.exists(jdk_dir):
        ensure_dir(jdk_dir)

        if label_type == "LAN":
            jdk_url = f"{server.rstrip('/')}{jdk_path}"
        else:
            jdk_url = f"{file_download_uri.rstrip('/')}{jdk_path}"

        headers = {"Redirect-Uri": file_download_uri}
        print("开始下载 JDK:", jdk_url)

        with tempfile.TemporaryDirectory() as td:
            tar_path = Path(td) / "jdk.tar.gz"
            http_download(jdk_url, tar_path, headers=headers, timeout=30000)
            extract_tar_gz(tar_path, jdk_dir)

    # 将解压目录重命名为 jdk1.8
    extracted_subdirs = [p for p in jdk_dir.iterdir() if p.is_dir() and p.name.startswith("jdk1.8")]
    target = jdk_dir / "jdk1.8"
    if extracted_subdirs:
        if target.exists():
            shutil.rmtree(target)
        extracted_subdirs[0].rename(target)
    else:
        # 回退：若未找到，构造两个常见目录名尝试
        for cname in ("jdk1.8.0_251", "jdk1.8.0_202"):
            cand = jdk_dir / cname
            if cand.exists():
                if target.exists():
                    shutil.rmtree(target)
                cand.rename(target)
                break

    java_bin = target / "bin" / "java"
    if not java_bin.exists():
        raise RuntimeError(f"JDK 安装失败：未找到 java 可执行文件，期望路径：{str(java_bin)}")

    print("JDK 安装完成:", target)
    return java_bin


def ensure_java(install_dir: Path, file_download_uri: str, server: str, label_type: str, jdk_path: str) -> Path:
    # 优先复用已安装在 install_dir 的内置 JDK
    jdk_base = install_dir / "jdk4agent"
    candidates = [
        jdk_base / "jdk1.8" / "bin" / "java",
        jdk_base / "jdk1.8.0_251" / "bin" / "java",
        jdk_base / "jdk1.8.0_202" / "bin" / "java",
    ]
    for local_java in candidates:
        if local_java.exists():
            ver = detect_java_version(str(local_java))
            print("内置 JDK 检测:", local_java, "版本:", ver)
            if ver and version_second_bit_ge_6(ver):
                return local_java

    # 若系统已存在且版本 >= 1.6 则可用
    system_java = which("java")
    if system_java:
        ver = detect_java_version(system_java)
        print("系统 Java:", system_java, "版本:", ver)
        if ver and version_second_bit_ge_6(ver):
            return Path(system_java)

    # 两者都不可用则下载内置 JDK
    print("未发现可用 Java，准备安装内置 JDK...")
    return install_jdk(install_dir, file_download_uri, server, label_type, jdk_path)


def agent_download(install_dir: Path, file_download_uri: str, server: str, label_type: str, *, force_download: bool = False) -> Path:
    if label_type == "LAN":
        url = f"{server.rstrip('/')}{AGENT_DOWNLOAD_PATH}"
    else:
        url = f"{file_download_uri.rstrip('/')}{AGENT_DOWNLOAD_PATH}"
    headers = {"Redirect-Uri": file_download_uri}
    dst = install_dir / "agent.jar"
    if dst.exists() and not force_download:
        print("检测到已存在 agent.jar，跳过下载。")
        return dst
    if dst.exists():
        dst.unlink()
    print("开始下载 agent:", url)
    http_download(url, dst, headers=headers, timeout=180)
    print("agent 下载完成:", dst)
    return dst


def proxy_download(install_dir: Path, file_download_uri: str, proxy_rel_path: str, *, force_download: bool = False) -> Path:
    url = f"{file_download_uri.rstrip('/')}{proxy_rel_path}"
    dst = install_dir / "gitee-proxy"
    if dst.exists() and not force_download:
        print("检测到已存在 gitee-proxy，跳过下载。")
        try:
            dst.chmod(0o755)
        except Exception:
            pass
        return dst
    if dst.exists():
        dst.unlink()
    print("开始下载 proxy:", url)
    http_download(url, dst, timeout=60)
    dst.chmod(0o755)
    print("proxy 下载完成:", dst)
    return dst


def get_local_ip() -> str:
    try:
        r = subprocess.run(["hostname", "-I"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, timeout=3)
        if r.stdout.strip():
            return r.stdout.strip().split()[0]
    except Exception:
        pass
    try:
        r = subprocess.run(["hostname", "-i"], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, timeout=3)
        if r.stdout.strip():
            return r.stdout.strip().split()[0]
    except Exception:
        pass
    return ""


def get_public_ip() -> str:
    try:
        data = http_get("https://ifconfig.co", timeout=3)
        return data.decode().strip()
    except Exception:
        return ""


def request_uuid_token(server: str, uuid_url_query: str, payload: dict) -> str | None:
    url = f"{server.rstrip('/')}{UUID_AUTH_URL_PATH}?{uuid_url_query}"
    data = json.dumps(payload).encode("utf-8")
    req = Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
    with urlopen(req, timeout=30) as resp:
        body = resp.read()
    try:
        obj = json.loads(body)
    except json.JSONDecodeError:
        print("服务端返回非 JSON：", body[:200], file=sys.stderr)
        return None
    # 原逻辑取 data 字段
    if isinstance(obj, dict) and "data" in obj:
        data_field = obj["data"]
        # data 可能是字符串或对象；若是对象且有 token/uuid 字段，尽量提取
        if isinstance(data_field, str):
            return data_field
        if isinstance(data_field, dict):
            for k in ("uuid", "token", "id"):
                if k in data_field and isinstance(data_field[k], str):
                    return data_field[k]
    return None


def read_existing_uuid(uuid_dir: Path) -> str | None:
    fp = uuid_dir / "uuid"
    if fp.exists():
        try:
            return fp.read_text(encoding="utf-8").strip()
        except Exception:
            return None
    return None


def write_uuid(uuid_dir: Path, token: str) -> None:
    ensure_dir(uuid_dir)
    (uuid_dir / "uuid").write_text(token + "\n", encoding="utf-8")
    tip = "This is the directory where the configuration file of Gitee Go Agent is located, please do not delete. //请勿删除\n"
    (uuid_dir / "Do Not Del This Directory").write_text(tip, encoding="utf-8")


def is_process_running(pattern: str) -> bool:
    # 使用 pgrep -f，若不可用则保守返回 False
    try:
        r = subprocess.run(["pgrep", "-f", pattern], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return r.returncode == 0
    except Exception:
        return False


def tail_file(path: Path, n: int = 200) -> str:
    try:
        with open(path, "rb") as f:
            f.seek(0, os.SEEK_END)
            size = f.tell()
            block = 4096
            data = bytearray()
            while len(data) < n * 120 and size > 0:
                read_size = min(block, size)
                size -= read_size
                f.seek(size)
                data[:0] = f.read(read_size)
            return b"".join(data.splitlines()[-n:]).decode(errors="ignore")
    except Exception:
        return ""


def launch_agent(java_bin: Path, install_dir: Path, server: str, token: str, force: bool = False) -> None:
    log_file = install_dir / "run.log"
    if log_file.exists():
        try:
            log_file.unlink()
        except Exception:
            pass
    running = is_process_running("agent.jar")
    if force or not running:
        print("启动 Agent...")
        with open(log_file, "ab", buffering=0) as lf:
            subprocess.Popen(
                [str(java_bin), "-Dfile.encoding=UTF-8", "-jar", "agent.jar", "-s", server, "-t", token],
                cwd=str(install_dir),
                stdout=lf,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
        time.sleep(10)
        print(tail_file(log_file))
    else:
        print("Agent 已在运行，无需重复启动。")


def launch_proxy(install_dir: Path, server: str, token: str) -> None:
    log_file = install_dir / "run.log"
    if log_file.exists():
        try:
            log_file.unlink()
        except Exception:
            pass
    running = is_process_running("gitee-proxy")
    if not running:
        print("启动 Proxy...")
        with open(log_file, "ab", buffering=0) as lf:
            subprocess.Popen(
                ["./gitee-proxy", "-s", server, "-t", token],
                cwd=str(install_dir),
                stdout=lf,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
        time.sleep(3)
    else:
        print("Proxy 已在运行，无需重复启动。")


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Gitee Go Agent 安装脚本（Python 版）")
    parser.add_argument("--server", "-s", default=DEFAULT_SERVER, help=f"Gitee Go 服务端地址 (默认: {DEFAULT_SERVER})")
    parser.add_argument("--uuid-url", "-u", required=False, help=("用于获取 Agent UUID 的认证 URL 参数（与 --reuse-uuid 二选一）。应当是一长串包括了labelId、sign、timestamp的字符串"))
    parser.add_argument("--remark", "-r", default="", help="Agent 备注信息")
    parser.add_argument("--count", "-c", default="", help="Agent 并发执行任务数")
    parser.add_argument("--proxy", "-p", dest="max_proxy", default="", help="作为 Proxy 时，可连接的 Agent 最大数量")
    parser.add_argument("--label", "-l", dest="label_type", default="", help="Agent 标签类型 (如: LAN)")
    parser.add_argument("--force-download", action="store_true", help="强制重新下载 JDK/agent/proxy（忽略已有文件）")
    parser.add_argument("--reuse-uuid", action="store_true", help="复用 ~/.gitee-agent/uuid，跳过服务端获取 token")
    parser.add_argument("--mode", choices=["agent", "proxy"], help="手动指定运行模式；未提供时从 --uuid-url 的查询参数解析")
    parser.add_argument(
        "--install-dir",
        default=str(Path.cwd() / "gitee_go_agent"),
        help=f"Agent 安装目录 (默认: {Path.cwd() / 'gitee_go_agent'})",
    )
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)

    arch, jdk_path, proxy_path = is_arch_supported()

    file_download_uri = DEFAULT_FILE_DOWNLOAD_URI
    install_dir = Path(args.install_dir)
    if not install_dir.is_absolute():
        install_dir = Path.cwd() / install_dir
    ensure_dir(install_dir)

    # 关键参数校验（与原脚本行为一致）
    if not file_download_uri:
        print("ERROR: file download uri is blank", file=sys.stderr)
        return 1
    if not jdk_path:
        print("ERROR: jdk download path is blank", file=sys.stderr)
        return 1
    if not AGENT_DOWNLOAD_PATH:
        print("ERROR: agent download path is blank", file=sys.stderr)
        return 1
    if not proxy_path:
        print("ERROR: proxy download path is blank", file=sys.stderr)
        return 1
    if not args.uuid_url and not args.reuse_uuid:
        print("ERROR: uuid url is blank (or use --reuse-uuid)", file=sys.stderr)
        return 1
    if not args.server:
        print("ERROR: server is blank", file=sys.stderr)
        return 1
    if not str(install_dir):
        print("ERROR: install-dir is blank", file=sys.stderr)
        return 1

    print("准备环境检查...")
    if not monitor_ok(args.server):
        print("无法访问服务端监控接口，请检查网络", file=sys.stderr)
        return 1

    # UUID 本地持久化目录
    uuid_dir = Path.home() / ".gitee-agent"

    # 先确定模式（优先命令行，其次从 uuid-url 解析，默认 agent）
    mode = (args.mode or "").strip()
    if not mode and args.uuid_url:
        q_for_mode = parse_qs(args.uuid_url, keep_blank_values=True)
        mode = (q_for_mode.get("mode", [""])[0] or "").strip()
    if not mode:
        mode = "agent"

    # 准备 token：复用或请求
    token: str | None = None
    if args.reuse_uuid:
        token = read_existing_uuid(uuid_dir)
        if not token:
            print("ERROR: 未找到可复用的 UUID（~/.gitee-agent/uuid），请提供 --uuid-url 或先完成一次安装", file=sys.stderr)
            return 1
        print("复用本地 UUID:", token)
    else:
        # 获取 IP 信息并从服务端请求 token
        local_ip = get_local_ip()
        public_ip = get_public_ip()
        payload = {
            "agentName": socket.gethostname(),
            "intranetIp": local_ip,
            "publicIp": public_ip,
            "count": str(args.count),
            "remark": str(args.remark),
            "maxProxyAgent": str(args.max_proxy),
        }
        print("请求 UUID token...")
        token = request_uuid_token(args.server, args.uuid_url or "", payload)
        print("UUID_TOKEN:", token)
        if not token or token == "null":
            print("未获取到有效的 UUID token", file=sys.stderr)
            return 1

    if mode == "proxy":
        print("进入 Proxy 模式：下载并启动代理...")
        proxy_download(install_dir, file_download_uri, proxy_path, force_download=args.force_download)
        write_uuid(uuid_dir, token)
        launch_proxy(install_dir, args.server, token)
        return 0

    elif mode == "agent":
        print("进入 Agent 模式：检查/安装 JDK 并下载 agent...")
        # 若本地已存在内置 JDK 且可用将复用；仅在必要时下载
        java_bin = ensure_java(install_dir, file_download_uri, args.server, args.label_type, jdk_path)
        # 若 agent.jar 已存在且未指定 --force-download 则跳过下载
        agent_download(install_dir, file_download_uri, args.server, args.label_type, force_download=args.force_download)

        existing = read_existing_uuid(uuid_dir)
        if existing and existing != token:
            # 与原逻辑一致：交互式确认
            print("-------------- Agent 安装提示 --------------")
            print("当前 Agent UUID:", existing)
            ans = input("当前主机环境已安装过 Agent ，继续执行将取消与之前主机组的绑定关系。是否继续? [Y/n] ").strip().lower()
            if ans in ("y", "yes", ""):
                print("将使用新的 UUID:", token)
                write_uuid(uuid_dir, token)
                launch_agent(java_bin, install_dir, args.server, token, force=True)
            elif ans in ("n", "no"):
                print("继续使用原有 UUID:", existing)
                write_uuid(uuid_dir, existing)
                launch_agent(java_bin, install_dir, args.server, existing, force=False)
            else:
                print("未选择有效选项，默认继续使用原有 UUID:", existing)
                write_uuid(uuid_dir, existing)
                launch_agent(java_bin, install_dir, args.server, existing, force=False)
        else:
            # 首次安装或相同 UUID
            write_uuid(uuid_dir, token)
            launch_agent(java_bin, install_dir, args.server, token, force=False)

        return 0

    else:
        raise RuntimeError(f"不支持的模式: {mode}")


if __name__ == "__main__":
    try:
        sys.exit(main(sys.argv[1:]))
    except KeyboardInterrupt:
        print("\n中断退出")
        sys.exit(130)
